import os
import numpy as np
import torch
from torch.utils.data import DataLoader
from pathlib import Path


class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, data_path):
        data = np.load(data_path, allow_pickle=True)
        self.images = torch.tensor(data['images'], dtype=torch.float32)
        self.labels = torch.tensor(data['labels'], dtype=torch.int64)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        return image, label


def load_client_data(is_train=True, datafile='data1', client_id=None):
    data_type = 'train' if is_train else 'test'

    # 默认路径：data/train/default_data
    default_data_path = Path(f'./data/{data_type}/default_data') / datafile
    # 客户端路径：data/train/{client_id}
    client_data_path = Path(f'./data/{data_type}/{client_id}') / datafile if client_id else None

    # 检查默认数据路径是否存在
    if default_data_path.exists():
        data_path = default_data_path
    elif client_data_path and client_data_path.exists():
        # 如果指定了client_id，并且在客户端路径中找到了文件
        data_path = client_data_path
    else:
        # 如果在两个路径下都没有找到文件
        raise FileNotFoundError(f"Data file '{datafile}' not found in either default or client directories.")

    # 加载数据集
    dataset = CustomDataset(data_path)

    return dataset
